#pragma once

#include <atomic>
#include <vector>

namespace hawking {
namespace indexlib {

class AtomicBitmap {
public:
    AtomicBitmap() = default;
    AtomicBitmap(size_t size) {
        size_ = (size + 31) / 32;
        bitmap_.reserve(size_);
        for (size_t idx = 0; idx < size_; ++idx) {
            bitmap_.emplace_back(std::make_unique<std::atomic<uint32_t>>(0));
        }
    }
    ~AtomicBitmap() = default;

    void Set(size_t value) {
        size_t elem_idx = value / 32;
        size_t bit_idx = value % 32;
        uint32_t mask = 1U << bit_idx;
        bitmap_[elem_idx]->fetch_or(mask, std::memory_order_release) & mask;
    }
    void Clear(size_t value) {
        size_t elem_idx = value / 32;
        size_t bit_idx = value % 32;
        uint32_t mask = ~(1U << bit_idx);
        bitmap_[elem_idx]->fetch_and(
            mask, std::memory_order_release) & (1U << bit_idx);
    }
    bool Test(size_t value) const {
        size_t elem_idx = value / 32;
        size_t bit_idx = value % 32;
        return bitmap_[elem_idx]->load(std::memory_order_release) & (1U << bit_idx);
    }
    size_t Count() const {
        size_t count = 0;
        for (const auto& elem : bitmap_) {
            count += __builtin_popcount(elem->load(std::memory_order_release));
        }
        return count;
    }

    template<class T>
    void Traverse(std::vector<T>* result) const {
        for (size_t idx = 0; idx < bitmap_.size(); ++idx) {
            uint32_t value = bitmap_[idx]->load(std::memory_order_acquire);
            for (auto bit_idx = 0; bit_idx < 32; ++bit_idx) {
                if (value & (1 << bit_idx)) {
                    result->emplace_back((idx << 5) + bit_idx);
                }
            }
        }
    }

    template<class T>
    void TraverseTopk(std::vector<T>* result, uint32_t topk) const {
        for (size_t idx = 0; idx < bitmap_.size(); ++idx) {
            uint32_t value = bitmap_[idx]->load(std::memory_order_acquire);
            for (auto bit_idx = 0; bit_idx < 32; ++bit_idx) {
                if (value & (1 << bit_idx)) {
                    result->emplace_back((idx << 5) + bit_idx);
                    if (result->size() == topk) {
                        return;
                    }
                }
            }
        }
    }

private:
    std::vector<std::unique_ptr<std::atomic<uint32_t>>> bitmap_;
    size_t size_;
};

}
}